﻿using edan.zhjk.Commons.DBOption;
using Furion;
using Microsoft.Extensions.Logging;
using ModelContextProtocol.Protocol;
using ModelContextProtocol.Server;
using Newtonsoft.Json;
using Npgsql;
using System.ComponentModel;
using System.Text;

namespace edan.zhjk.Commons.MCPService.PostgreSQLDB
{
    /// <summary>
    /// pgsql数据库Mcp工具类
    /// </summary>
    [McpServerToolType, Description("PostgreSQL数据库")]
    public sealed class PostgreSQLDBTool
    {
        private readonly ILogger<PostgreSQLDBTool> _logger;
        private static MutiDBConns? MutiDBConns;
        /// <summary>
        /// 构造函数
        /// </summary>
        /// <param name="logger"></param>
        public PostgreSQLDBTool(ILogger<PostgreSQLDBTool> logger)
        {
            _logger = logger;
            var dbOptions = App.GetOptionsMonitor<DBOptions>();
            MutiDBConns = dbOptions?.MutiDBConns?.FirstOrDefault(it => it.DbType == DataBaseType.PostgreSQL);
        }

        /// <summary>
        /// 读取PGSql连接
        /// </summary>
        /// <param name="url">数据库链接地址</param>
        /// <returns></returns>
        /// <exception cref="InvalidOperationException"></exception>
        private NpgsqlConnection GetNpgsqlConnection(string? url)
        {
            var finalUrl = string.IsNullOrWhiteSpace(url) ? MutiDBConns?.Connection : url;

            if (string.IsNullOrWhiteSpace(finalUrl))
                throw new InvalidOperationException("数据库链接地址字段是必需的，但未在环境中提供或设置。");

            return new NpgsqlConnection(finalUrl);
        }

        /// <summary>
        /// 获取数据库中的所有表
        /// </summary>
        /// <param name="url"></param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        [McpServerTool(Name = "postgres_get_tables"), Description("[PostgreSQL]获取数据库中的所有表")]
        public async Task<CallToolResult> GetTables([Description("数据库链接字符串")] string? url = null
            , CancellationToken cancellationToken = default)
        {
            try
            {
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                await using var cmd = conn.CreateCommand();
                cmd.CommandText = @"
SELECT t.schemaname AS ""表所在的模式名"", t.tablename AS ""表名"", t.tableowner AS ""表的所有者""
     , t.tablespace AS ""表所在的表空间"", t.hasindexes AS ""表是否有索引"", t.hasrules AS ""表是否有规则""
     , t.hastriggers AS ""表是否有触发器"", t.rowsecurity AS ""表是否启用行级安全策略""
FROM pg_tables t
WHERE t.schemaname='public';";
                await using var reader = await cmd.ExecuteReaderAsync(cancellationToken);

                var schemas = new List<string>();
                while (await reader.ReadAsync(cancellationToken))
                {
                    schemas.Add(reader.GetString(0));
                }

                var settings = new JsonSerializerSettings
                {
                    Formatting = Formatting.None,
                    NullValueHandling = NullValueHandling.Include
                };
                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = JsonConvert.SerializeObject(schemas, settings) }]
                };
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = ex.Message }]
                };
            }
        }

        /// <summary>
        /// 查询所有表结构明细
        /// </summary>
        /// <param name="url"></param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        [McpServerTool(Name = "postgres_get_table_schemas"), Description("[PostgreSQL]获取数据库所有表结构明细")]
        public async Task<CallToolResult> GetTableSchemas([Description("数据库链接字符串")] string? url = null
            , CancellationToken cancellationToken = default)
        {
            try
            {
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                await using var cmd = conn.CreateCommand();
                cmd.CommandText = @"
SELECT pa.attnum AS ""序号""
  ,pc.relname AS ""表名""
  ,cast(obj_description(relfilenode,'pg_class') as varchar) AS ""表名描述""
  ,pa.attname AS ""列名""
  ,concat_ws('' ,pt.typname ,SUBSTRING(format_type(pa.atttypid ,pa.atttypmod) FROM '\(.*\)')) AS ""字段类型""
  ,pa.attnotnull AS ""必填""
  ,(CASE
		WHEN (SELECT count(pg_constraint.*) FROM pg_constraint
		INNER JOIN pg_class ON pg_constraint.conrelid = pg_class.oid
		INNER JOIN pg_attribute ON pg_attribute.attrelid = pg_class.oid
			AND pg_attribute.attnum = any(pg_constraint.conkey)
		INNER JOIN pg_type ON pg_type.oid = pg_attribute.atttypid
		WHERE pg_class.relname = pc.relname
			AND pg_constraint.contype = 'p'
			AND pg_attribute.attname = pa.attname) > 0 THEN TRUE
		ELSE FALSE END
  ) AS ""主键""
  ,pd.description AS ""备注""
FROM pg_class pc, pg_attribute pa , pg_type pt, pg_description pd
WHERE pa.attnum>0
  AND pa.attrelid = pc.oid
  AND pa.atttypid = pt.oid
  AND pd.objoid=pa.attrelid
  AND pd.objsubid=pa.attnum
ORDER BY pc.relname DESC ,pa.attnum ASC;";
                await using var reader = await cmd.ExecuteReaderAsync(cancellationToken);

                var schemas = new List<string>();
                while (await reader.ReadAsync(cancellationToken))
                {
                    schemas.Add(reader.GetString(0));
                }

                var settings = new JsonSerializerSettings
                {
                    Formatting = Formatting.None,
                    NullValueHandling = NullValueHandling.Include
                };
                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = JsonConvert.SerializeObject(schemas, settings) }]
                };
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = ex.Message }]
                };
            }
        }

        /// <summary>
        /// 检索并返回有关表的完整元数据，包括列、主键和外键。
        /// </summary>
        /// <param name="schema">模式名称</param>
        /// <param name="table">表名</param>
        /// <param name="url">数据库链接字符串</param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        [McpServerTool(Name = "postgres_describe_table"), Description("[PostgreSQL]获取数据库有关表的完整元数据，包括列、主键和外键。")]
        public async Task<CallToolResult> DescribeTable(
            [Description("模式名称")] string? schema = null,
            [Description("表名")] string table = "",
            [Description("数据库链接字符串")] string? url = null,
            CancellationToken cancellationToken = default)
        {
            try
            {
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                await using var cmd = conn.CreateCommand();
                cmd.CommandText = $@"SELECT * FROM information_schema.COLUMNS WHERE table_schema = '{schema}' AND TABLE_NAME = '{table}';";
                await using var reader = await cmd.ExecuteReaderAsync(cancellationToken);

                var schemas = new List<string>();
                while (await reader.ReadAsync(cancellationToken))
                {
                    schemas.Add(reader.GetString(0));
                }

                var settings = new JsonSerializerSettings
                {
                    Formatting = Formatting.None,
                    NullValueHandling = NullValueHandling.Include
                };
                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = JsonConvert.SerializeObject(schemas, settings) }]
                };
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = ex.Message }]
                };
            }
        }

        /// <summary>
        /// 执行一个 SQL 查询，并以 JSONL 格式返回结果。
        /// </summary>
        /// <param name="query"></param>
        /// <param name="max_rows"></param>
        /// <param name="url"></param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        [McpServerTool(Name = "postgres_execute_query"), Description("[PostgreSQL]执行一个SQL查询，并以JSONL格式返回结果。")]
        public async Task<CallToolResult> ExecuteQuery(
            [Description("SQL查询条件")] string query,
            [Description("最大数据条数")] int? max_rows = null,
            [Description("数据库链接字符串")] string? url = null,
            CancellationToken cancellationToken = default)
        {
            var maxRows = max_rows ?? 100;
            try
            {
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                await using var command = conn.CreateCommand();
                command.CommandText = query;
                await using var reader = await command.ExecuteReaderAsync(cancellationToken);

                var list = new List<Dictionary<string, object?>>();
                var settings = new JsonSerializerSettings
                {
                    Formatting = Formatting.None,
                    NullValueHandling = NullValueHandling.Include
                };

                int count = 0;
                while (await reader.ReadAsync(cancellationToken) && count < maxRows)
                {
                    var row = new Dictionary<string, object?>();
                    for (int i = 0; i < reader.FieldCount; i++)
                    {
                        var val = reader.GetValue(i);
                        var str = val?.ToString() ?? null;
                        if (str != null && str.Length > MutiDBConns?.MaxLongData)
                            str = str.Substring(0, MutiDBConns?.MaxLongData ?? 100);
                        row[reader.GetName(i)] = str;
                    }
                    list.Add(row);
                    count++;
                }
                var text = JsonConvert.SerializeObject(list, settings);
                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = text }]
                };
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = ex.Message }]
                };
            }
        }

        /// <summary>
        /// 执行一个 SQL 查询，并以 Markdown 表格格式返回结果。
        /// </summary>
        /// <param name="query"></param>
        /// <param name="max_rows"></param>
        /// <param name="url">数据库链接字符串</param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        [McpServerTool(Name = "postgres_execute_query_md"), Description("[PostgreSQL]执行一个SQL查询，并以 Markdown 表格格式返回结果。")]
        public async Task<CallToolResult> ExecuteQueryMd(
            [Description("SQL查询条件")] string query,
            [Description("最大数据条数")] int? max_rows = null,
            [Description("数据库链接字符串")] string? url = null,
            CancellationToken cancellationToken = default)
        {
            var maxRows = max_rows ?? 100;
            try
            {
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                await using var cmd = conn.CreateCommand();
                cmd.CommandText = query;
                await using var reader = await cmd.ExecuteReaderAsync(cancellationToken);

                var colNames = Enumerable.Range(0, reader.FieldCount)
                    .Select(i => reader.GetName(i))
                    .ToList();

                var md = new StringBuilder();
                md.Append("| ").Append(string.Join(" | ", colNames)).AppendLine(" |");
                md.Append("| ").Append(string.Join(" | ", colNames.Select(n => "---"))).AppendLine(" |");

                int count = 0;
                while (await reader.ReadAsync(cancellationToken) && count < maxRows)
                {
                    md.Append("| ");
                    for (int i = 0; i < reader.FieldCount; i++)
                    {
                        var val = reader.GetValue(i)?.ToString() ?? string.Empty;
                        if (val.Length > MutiDBConns?.MaxLongData)
                            val = val.Substring(0, MutiDBConns?.MaxLongData ?? 100);
                        md.Append(val).Append(" | ");
                    }
                    md.AppendLine();
                    count++;
                }
                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = md.ToString() }]
                };
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = ex.Message }]
                };
            }
        }

        /// <summary>
        /// 执行一个 SQL 查询，并以 JSONL 格式返回结果。
        /// </summary>
        /// <param name="query"></param>
        /// <param name="url">数据库链接字符串</param>
        /// <param name="cancellationToken"></param>
        /// <returns></returns>
        /// <exception cref="ArgumentException"></exception>
        [McpServerTool(Name = "postgres_query_database"), Description("[PostgreSQL]执行一个SQL查询，并以 JSONL 格式返回结果。")]
        public async Task<CallToolResult> QueryDatabase(
            [Description("SQL查询条件")] string query,
            [Description("数据库链接字符串")] string? url = null,
            CancellationToken cancellationToken = default)
        {
            try
            {
                if (string.IsNullOrWhiteSpace(query))
                    throw new ArgumentException("Query string cannot be null or empty.", nameof(query));
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                await using var cmd = conn.CreateCommand();
                cmd.CommandText = query;
                await using var reader = await cmd.ExecuteReaderAsync(cancellationToken);

                var list = new List<Dictionary<string, object?>>();
                var settings = new JsonSerializerSettings
                {
                    Formatting = Formatting.None,
                    NullValueHandling = NullValueHandling.Include
                };

                while (await reader.ReadAsync(cancellationToken))
                {
                    var row = new Dictionary<string, object?>();
                    for (int i = 0; i < reader.FieldCount; i++)
                    {
                        var val = reader.GetValue(i);
                        var str = val?.ToString() ?? null;
                        if (str != null && str.Length > MutiDBConns?.MaxLongData)
                            str = str.Substring(0, MutiDBConns?.MaxLongData ?? 100);
                        row[reader.GetName(i)] = str;
                    }
                    list.Add(row);
                }
                var text = JsonConvert.SerializeObject(list, settings);
                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = text }]
                };
            }
            catch (Exception ex) when (ex is not OperationCanceledException)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = ex.Message }]
                };
            }
        }

        /// <summary>
        /// 这是供人工智能测试数据库连接的工具
        /// </summary>
        /// <returns></returns>
        [McpServerTool(Name = "postgres_health_check"), Description("[PostgreSQL]检查数据库连接是否正常且可用")]
        public async Task<CallToolResult> HealthCheck([Description("数据库链接字符串")] string? url = null
            , CancellationToken cancellationToken = default)
        {
            try
            {
                await using var conn = GetNpgsqlConnection(url);
                await conn.OpenAsync(cancellationToken);

                return new CallToolResult
                {
                    IsError = false,
                    Content = [new TextContentBlock() { Type = "text", Text = "true" }]
                };
            }
            catch (Exception e)
            {
                return new CallToolResult
                {
                    IsError = true,
                    Content = [new TextContentBlock() { Type = "text", Text = "false" }]
                };
            }
        }
    }
}